package zttc.itat.user.util;

import static org.junit.Assert.assertNotNull;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileWriter;
import java.io.IOException;
import java.io.InputStream;
import java.sql.SQLException;

import org.dbunit.DatabaseUnitException;
import org.dbunit.database.DatabaseConnection;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.database.QueryDataSet;
import org.dbunit.dataset.DataSetException;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSet;
import org.dbunit.operation.DatabaseOperation;
import org.junit.AfterClass;
import org.junit.BeforeClass;

public class AbstractDbUtilTestCase {
	public static IDatabaseConnection dbunitConn;
	private File tempFile;
	
	@BeforeClass
	public static void init() throws ClassNotFoundException, DatabaseUnitException, SQLException{
		dbunitConn=new DatabaseConnection(DbUtil.getConnection());
	}
	
	
	protected IDataSet createDataSet(String tname) throws DataSetException, IOException{
		InputStream is = AbstractDbUtilTestCase.class.getClassLoader().getResourceAsStream(tname+".xml");
		assertNotNull("dbunit的基础数据文件不存在",is);
		return new FlatXmlDataSet(is);
	}
	
	protected void backupAllTable() throws SQLException, DataSetException, IOException{
		IDataSet ds = dbunitConn.createDataSet();	
		writeBackupFile(ds);
	}
	
	private void writeBackupFile(IDataSet ds) throws IOException, DataSetException{
		tempFile=File.createTempFile("back", "xml");
		FlatXmlDataSet.write(ds, new FileWriter(tempFile));
	}
	
	protected void backupCustomTable(String[] tname) throws DataSetException, IOException, SQLException{
		QueryDataSet ds = new QueryDataSet(dbunitConn);
		for(String str:tname){
			ds.addTable(str);
		}
		writeBackupFile(ds);
	}
	
	protected void backupOneTable(String tname) throws DataSetException, IOException, SQLException{
		backupCustomTable(new String[]{tname});
	}
	
	protected void resumeTable() throws DatabaseUnitException, SQLException, IOException{
		InputStream is = new FileInputStream(tempFile);
		IDataSet ds = new FlatXmlDataSet(tempFile);
		DatabaseOperation.CLEAN_INSERT.execute(dbunitConn, ds);
	}
	
	
	@AfterClass
	public static void destroy(){
		if(dbunitConn!=null){
			try {
				dbunitConn.close();
			} catch (SQLException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			}
		}
	}
}
